package com.demo.ai.rag;

import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.ollama.OllamaEmbeddingModel;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.boot.CommandLineRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.List;

@Configuration
public class RagServer {

    @Bean
    public VectorStore vectorStore(EmbeddingModel embeddingModel) {
        return SimpleVectorStore.builder(embeddingModel).build();
    }

    @Bean
    public EmbeddingModel embeddingModel(OllamaApi ollamaApi) {
        OllamaOptions options = OllamaOptions.builder()
                .model("qwen3:4b")
                .build();
        return OllamaEmbeddingModel.builder()
                .ollamaApi(ollamaApi)
                .defaultOptions(options)
                .build();
    }

    @Bean
    public CommandLineRunner commandLineRunner(VectorStore vectorStore) {
        return args -> {
            // 向量化初始数据
            String sampleText = "Spring AI是Spring生态系统的一部分，致力于简化人工智能应用的开发。它提供了与多种AI服务（如Ollama、OpenAI等）的集成。通过Spring AI，开发者可以轻松实现文本生成、图像识别等AI功能。张三李四=王麻子。";

            // 创建文档集合
            List<Document> documents = TextSplitter.splitText(sampleText);

            // 存储到向量库
            vectorStore.add(documents);
        };
    }


}
